#include "PerspectiveCamera.cuh"

using namespace OBR;

extern "C"
{
    __constant__ Params params;
}

const float jitter_scale = 1.0;

extern "C" __global__ void __raygen__perspective_camera()
{
    RaygenData *sbt_record = reinterpret_cast<RaygenData *>(optixGetSbtDataPointer());
    PerspectiveCameraData *camera_data =
        reinterpret_cast<PerspectiveCameraData *>(sbt_record->camera_payload);
    const uint3 launch_index = optixGetLaunchIndex();
    const uint3 launch_dims = optixGetLaunchDimensions();
    const uint2 pos = make_uint2(launch_index.x, launch_index.y);
    const uint image_index = launch_index.y * params.width + launch_index.x;
    bool log = SHADER_LOG_COND;

#ifdef DEBUG_SHADER
    if (log)
    {
        printf("Camera data:\n");
        printf("orig: ");
        print_float3(camera_data->orig);
        printf("\n");

        printf("U: ");
        print_float3(camera_data->U);
        printf("\n");

        printf("V: ");
        print_float3(camera_data->V);
        printf("\n");

        printf("W: ");
        print_float3(camera_data->W);
        printf("\n");
    }
#endif

    // init RNG
    curandState_t rng;
    curand_init(launch_index.x, launch_index.y, params.launch_time, &rng);
    curandStatePhilox4_32_10_t rng_long;
    curand_init(launch_index.y, launch_index.x, params.launch_time, &rng_long);

    // init prd
    SamplePRD prd;
    prd.params = &params;
    prd.launch_index = launch_index;
    prd.log = log;
    prd.rng = &rng;
    prd.rng_long = &rng_long;

    Spectrum final_radiance;
    for (size_t s = 0; s < params.samples_per_launch; s++)
    {

#ifdef DEBUG_SHADER
        if (log)
        {
            printf("sample %d\n", s);
        }
#endif

        // generate ray
        float2 subpixel_jitter = make_float2(curand_uniform(&rng), curand_uniform(&rng)) - 0.5;
        subpixel_jitter *= jitter_scale;
        float2 inv_filmsize = 1.0f / make_float2(params.width, params.height);
        float2 d = make_float2(static_cast<float>(launch_index.x) + 0.5f + subpixel_jitter.x,
                               static_cast<float>(launch_index.y) + 0.5f + subpixel_jitter.y);
        RayData ray;
        ray.o = camera_data->orig;
        ray.d = normalize(d.x * inv_filmsize.x * camera_data->U +
                          d.y * inv_filmsize.y * camera_data->V + camera_data->W);
        ray.med_id = sbt_record->med_id;
        ray.med_instance_id = sbt_record->med_instance_id;
        ray.has_differential = true;
        ray.rx_o = ray.o;
        ray.rx_d = normalize((d.x + 1.0f) * inv_filmsize.x * camera_data->U +
                             d.y * inv_filmsize.y * camera_data->V + camera_data->W);
        ray.ry_o = ray.o;
        ray.ry_d = normalize(d.x * inv_filmsize.x * camera_data->U +
                             (d.y + 1.0f) * inv_filmsize.y * camera_data->V + camera_data->W);

#ifdef DEBUG_SHADER
        if (log)
        {
            printf("generated ray orig(%f, %f, %f), dir(%f, %f, %f)\n", ray.o.x, ray.o.y, ray.o.z,
                   ray.d.x, ray.d.y, ray.d.z);
        }
#endif

        // set sample specified params
        prd.scene_time = curand_uniform(&rng) * (params.tmax - params.tmin) + params.tmin;
        prd.ray_occlusion_test = false;
        prd.inter.wi = ray;
        // call integrator
        final_radiance += optixContinuationCall<Spectrum, SamplePRD *>(0, &prd);
    }

    // set radiance to fbo #0
    set_fbo(&params, launch_index, 0, final_radiance / params.samples_per_launch);
}